package com.pdwbdn.utils;

import org.apache.commons.codec.binary.Base64;

import javax.crypto.Cipher;
import java.io.*;
import java.nio.charset.StandardCharsets;
import java.security.*;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;

public class RSAUtil {

    /**
     * 签名算法
     */
    public static final String SIGNATURE_ALGORITHM = "MD5withRSA";
    /**
     * KEY算法
     */
    public static final String KEY_ALGORITHM = "RSA";

    /**
     * 随机生成密钥对
     */
    public static void genKeyPair(String outputPath) {
        if (outputPath == null) {
            return;
        }
        FileWriter publicFileWriter = null;
        BufferedWriter publicBufferedWriter = null;
        FileWriter privateFileWriter = null;
        BufferedWriter privateBufferedWriter = null;
        try {
            // KeyPairGenerator类用于生成公钥和私钥对，基于RSA算法生成对象
            KeyPairGenerator keyPairGen = KeyPairGenerator.getInstance(KEY_ALGORITHM);
            // 初始化密钥对生成器，密钥大小至少1024位
            keyPairGen.initialize(4096, new SecureRandom());
            // 生成一个密钥对，保存在keyPair中
            KeyPair keyPair = keyPairGen.generateKeyPair();
            // 得到私钥
            RSAPrivateKey privateKey = (RSAPrivateKey) keyPair.getPrivate();
            // 得到公钥
            RSAPublicKey publicKey = (RSAPublicKey) keyPair.getPublic();
            // 得到公钥字符串
            byte[] publicKeyBytes = Base64.encodeBase64Chunked(publicKey.getEncoded());
            String publicKeyStr = new String(publicKeyBytes, StandardCharsets.UTF_8);
            // 得到私钥字符串
            byte[] privateKeyBytes = Base64.encodeBase64Chunked(privateKey.getEncoded());
            String privateKeyStr = new String(privateKeyBytes, StandardCharsets.UTF_8);
            // 将密钥对写入到文件
            publicFileWriter = new FileWriter(outputPath + "/publickey.pem");
            publicBufferedWriter = new BufferedWriter(publicFileWriter);
            privateFileWriter = new FileWriter(outputPath + "/privatekey.pem");
            privateBufferedWriter = new BufferedWriter(privateFileWriter);
            // 写公钥
            publicBufferedWriter.write("-----BEGIN PUBLIC KEY-----\r\n");
            publicBufferedWriter.write(publicKeyStr);
            publicBufferedWriter.write("-----END PUBLIC KEY-----\r\n");
            publicBufferedWriter.flush();
            // 写私钥
            privateBufferedWriter.write("-----BEGIN PRIVATE KEY-----\r\n");
            privateBufferedWriter.write(privateKeyStr);
            privateBufferedWriter.write("-----END PRIVATE KEY-----\r\n");
            privateBufferedWriter.flush();
        } catch (Throwable t) {
            System.err.println(t.toString());
        } finally {
            // 关闭写入
            if (publicBufferedWriter != null) {
                try {
                    publicBufferedWriter.close();
                } catch (Throwable t) {
                    System.err.println(t.toString());
                }
            }
            if (publicFileWriter != null) {
                try {
                    publicFileWriter.close();
                } catch (Throwable t) {
                    System.err.println(t.toString());
                }
            }
            if (privateBufferedWriter != null) {
                try {
                    privateBufferedWriter.close();
                } catch (Throwable t) {
                    System.err.println(t.toString());
                }
            }
            if (privateFileWriter != null) {
                try {
                    privateFileWriter.close();
                } catch (Throwable t) {
                    System.err.println(t.toString());
                }
            }
        }
    }

    /**
     * 从文件中读取密钥字符串
     */
    public static String readKeyStrFromFile(String filePath) {
        if (filePath == null) {
            return null;
        }
        FileReader fileReader = null;
        BufferedReader bufferedReader = null;
        try {
            File file = new File(filePath);
            if (!file.exists() || !file.isFile()) {
                return null;
            }
            fileReader = new FileReader(file);
            bufferedReader = new BufferedReader(fileReader);
            String line;
            StringBuilder sb = new StringBuilder();
            while ((line = bufferedReader.readLine()) != null) {
                if (line.startsWith("---")) {
                    continue;
                }
                sb.append(line);
            }
            return sb.toString();
        } catch (Throwable t) {
            System.err.println(t.toString());
            return null;
        } finally {
            if (bufferedReader != null) {
                try {
                    bufferedReader.close();
                } catch (Throwable t) {
                    System.err.println(t.toString());
                }
            }
            if (fileReader != null) {
                try {
                    fileReader.close();
                } catch (Throwable t) {
                    System.err.println(t.toString());
                }
            }
        }
    }

    /**
     * 从字符串中加载公钥
     */
    public static RSAPublicKey loadPublicKeyByStr(String publicKeyStr) {
        if (publicKeyStr == null || publicKeyStr.isEmpty()) {
            return null;
        }
        try {
            String base64String = publicKeyStr.replaceAll("\\s", "");
            byte[] buffer = Base64.decodeBase64(base64String);
            if (buffer == null) {
                return null;
            }
            KeyFactory keyFactory = KeyFactory.getInstance(KEY_ALGORITHM);
            X509EncodedKeySpec keySpec = new X509EncodedKeySpec(buffer);
            return (RSAPublicKey) keyFactory.generatePublic(keySpec);
        } catch (Throwable t) {
            System.err.println(t.toString());
            return null;
        }
    }

    /**
     * 从字符串中加载私钥
     */
    public static RSAPrivateKey loadPrivateKeyByStr(String privateKeyStr) {
        if (privateKeyStr == null || privateKeyStr.isEmpty()) {
            return null;
        }
        try {
            String base64String = privateKeyStr.replaceAll("\\s", "");
            byte[] buffer = Base64.decodeBase64(base64String);
            if (buffer == null) {
                return null;
            }
            PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(buffer);
            KeyFactory keyFactory = KeyFactory.getInstance(KEY_ALGORITHM);
            return (RSAPrivateKey) keyFactory.generatePrivate(keySpec);
        } catch (Throwable t) {
            System.err.println(t.toString());
            return null;
        }
    }

    /**
     * 使用公钥加密
     */
    public static byte[] encrypt(RSAPublicKey publicKey, byte[] plainTextData) {
        if (publicKey == null || plainTextData == null || plainTextData.length < 1) {
            return null;
        }
        try {
            // 使用默认RSA
            Cipher cipher = Cipher.getInstance(KEY_ALGORITHM);
            cipher.init(Cipher.ENCRYPT_MODE, publicKey);
            return cipher.doFinal(plainTextData);
        } catch (Throwable t) {
            System.err.println(t.toString());
            return null;
        }
    }

    /**
     * 使用私钥加密
     */
    public static byte[] encrypt(RSAPrivateKey privateKey, byte[] plainTextData) {
        if (privateKey == null || plainTextData == null || plainTextData.length < 1) {
            return null;
        }
        try {
            // 使用默认RSA
            Cipher cipher = Cipher.getInstance(KEY_ALGORITHM);
            cipher.init(Cipher.ENCRYPT_MODE, privateKey);
            return cipher.doFinal(plainTextData);
        } catch (Throwable t) {
            System.err.println(t.toString());
            return null;
        }
    }

    /**
     * 使用私钥解密
     */
    public static byte[] decrypt(RSAPrivateKey privateKey, byte[] cipherData) {
        if (privateKey == null || cipherData == null || cipherData.length < 1) {
            return null;
        }
        try {
            // 使用默认RSA
            Cipher cipher = Cipher.getInstance(KEY_ALGORITHM);
            cipher.init(Cipher.DECRYPT_MODE, privateKey);
            return cipher.doFinal(cipherData);
        } catch (Throwable t) {
            System.err.println(t.toString());
            return null;
        }
    }

    /**
     * 使用公钥解密
     */
    public static byte[] decrypt(RSAPublicKey publicKey, byte[] cipherData) {
        if (publicKey == null || cipherData == null || cipherData.length < 1) {
            return null;
        }
        try {
            // 使用默认RSA
            Cipher cipher = Cipher.getInstance(KEY_ALGORITHM);
            cipher.init(Cipher.DECRYPT_MODE, publicKey);
            return cipher.doFinal(cipherData);
        } catch (Throwable t) {
            System.err.println(t.toString());
            return null;
        }
    }

    /**
     * 私钥签名
     */
    public static String sign(PrivateKey privateKey, byte[] data) {
        if (privateKey == null || data == null || data.length < 1) {
            return null;
        }
        try {
            Signature signature = Signature.getInstance(SIGNATURE_ALGORITHM);
            signature.initSign(privateKey);
            signature.update(data);
            byte[] signBytes = signature.sign();
            return Base64.encodeBase64String(signBytes);
        } catch (Throwable t) {
            System.err.println(t.toString());
            return null;
        }
    }

    /**
     * 公钥校验
     */
    public static boolean verify(PublicKey publicKey, byte[] data, String sign) {
        if (publicKey == null || data == null || data.length < 1 || sign == null || sign.isEmpty()) {
            return false;
        }
        try {
            Signature signature = Signature.getInstance(SIGNATURE_ALGORITHM);
            signature.initVerify(publicKey);
            signature.update(data);
            String base64String = sign.replaceAll("\\s", "");
            byte[] signBytes = Base64.decodeBase64(base64String);
            return signature.verify(signBytes);
        } catch (Throwable t) {
            System.err.println(t.toString());
            return false;
        }
    }
}

